#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
简化版CLIP模型演示
这是一个最简单的CLIP使用示例，适合初学者理解
"""

import torch
import clip
from PIL import Image
import numpy as np

def simple_clip_demo():
    """
    最简单的CLIP演示
    """
    print("=== 简单CLIP演示 ===")
    
    # 1. 加载模型
    print("1. 加载CLIP模型...")
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model, preprocess = clip.load("ViT-B/32", device=device)
    print(f"   使用设备: {device}")
    
    # 2. 创建一个简单的测试图像（红色圆圈）
    print("\n2. 创建测试图像...")
    img_array = np.zeros((224, 224, 3), dtype=np.uint8)
    
    # 画一个红色圆圈
    center = 112
    radius = 50
    y, x = np.ogrid[:224, :224]
    mask = (x - center) ** 2 + (y - center) ** 2 <= radius ** 2
    img_array[mask] = [255, 0, 0]  # 红色
    
    image = Image.fromarray(img_array)
    print("   创建了一个红色圆圈图像")
    
    # 3. 准备图像和文本
    image_input = preprocess(image).unsqueeze(0).to(device)
    
    # 候选文本描述
    texts = [
        "a red circle",
        "a blue square", 
        "a green triangle",
        "a yellow star"
    ]
    
    text_inputs = clip.tokenize(texts).to(device)
    
    # 4. 进行预测
    print("\n3. 进行图像-文本匹配...")
    with torch.no_grad():
        logits_per_image, logits_per_text = model(image_input, text_inputs)
        probs = logits_per_image.softmax(dim=-1).cpu().numpy()
    
    # 5. 显示结果
    print("\n4. 匹配结果:")
    print("-" * 30)
    
    for i, text in enumerate(texts):
        probability = probs[0][i] * 100
        print(f"   {text:<20} {probability:.1f}%")
    
    # 找出最匹配的描述
    best_match_idx = np.argmax(probs[0])
    best_match = texts[best_match_idx]
    best_prob = probs[0][best_match_idx] * 100
    
    print(f"\n最佳匹配: '{best_match}' (概率: {best_prob:.1f}%)")
    
    print("\n=== 演示完成 ===")
    print("\nCLIP模型成功识别出图像中的红色圆圈！")
    print("这展示了CLIP理解图像和文本关系的能力。")

def main():
    try:
        simple_clip_demo()
    except ImportError:
        print("错误: 缺少必要的依赖包")
        print("请运行: pip install torch clip-by-openai pillow numpy")
    except Exception as e:
        print(f"运行错误: {e}")

if __name__ == "__main__":
    main()